# 根据约束条件创建业务量矩阵
from copy import deepcopy
import math
import random
from network import QKDNetwork
from docplex.mp.model import Model
import numpy as np

m = []

def get_lp_solve(n : QKDNetwork):
    numNodes = len(n.nodeList) #节点的个数
    model = Model(name='LP_Matrix')
    totalCapacity = 0
    for i in range(numNodes):
        nodeID = i + 1 # --> i
        for j in n.getNode(nodeID).adjacentNodes:
            endNodeID = j
            totalCapacity += n.getLinkCapacity(nodeID, endNodeID)
    d = {}
    tmp_beta = model.quad_expr()
    for i in range(numNodes + 1):
        d[i] = {}
        for j in range(numNodes + 1):
            d[i][j] = model.integer_var(name=f"d_{i}_{j}", lb=0)
    d_ret = {}
    for i in range(numNodes):
        d_ret[i] = {}
        for j in range(numNodes):
            d_ret[i][j] = 0
    for i in range(numNodes):
        nodeID = i + 1 # --> i
        tmp_start_flow = model.linear_expr()
        tmp_end_flow = model.linear_expr()
        for j in range(numNodes):
            endNodeID = j + 1 # --> j
            if (nodeID != endNodeID):
                tmp_start_flow += d[nodeID][endNodeID]
        for j in n.getNode(nodeID).adjacentNodes:
            endNodeID = j
            tmp_start_flow -= n.getLinkCapacity(nodeID, endNodeID)
        for j in range(numNodes):
            endNodeID = j + 1 # --> j
            if (nodeID != endNodeID):
                tmp_end_flow += d[endNodeID][nodeID]
        for k in n.getNode(nodeID).adjacentNodes:
            endNodeID = k # --> k
            tmp_end_flow -= n.getLinkCapacity(endNodeID, nodeID)
        model.add_constraint(tmp_start_flow <= 0)
        model.add_constraint(tmp_end_flow <= 0)
    for i in range(numNodes):
        for j in range(numNodes):
            if (i != j):
                tmp_beta += (d[i + 1][j + 1] - totalCapacity) ** 2
    beta = model.continuous_var(name="Beta")
    beta = tmp_beta
    model.minimize(beta)
    # model.export_as_lp(f"lp/model_matrix.lp")
    sol = model.solve()
    if sol:
        for i in range(numNodes + 1):
            for j in range(numNodes + 1):
                # 随机挑选一些SD对
                if (i > 0 and j > 0):
                    d_ret[i - 1][j - 1] = math.ceil(sol[d[i][j]])
        return d_ret
    

def get_matrix_with_mean_capacity(nodeNum = 15):
    mean_capacity = 30
    matrix = np.zeros((nodeNum, nodeNum))
    non_diag_elements = set()
    while len(non_diag_elements) < nodeNum * (nodeNum - 1):
        random_value = np.random.uniform(1, 60)  # 选择一个范围来生成随机数
        if random_value not in non_diag_elements:
            non_diag_elements.add(random_value)
    non_diag_elements = list(non_diag_elements)
    np.random.shuffle(non_diag_elements)
    index = 0
    for i in range(nodeNum):
        for j in range(nodeNum):
            if i != j:
                matrix[i, j] = non_diag_elements[index]
                index += 1
    current_mean = np.mean(matrix[matrix != 0])
    # 均值调整为30
    adjustment_factor = mean_capacity / current_mean
    matrix[matrix != 0] *= adjustment_factor
    return matrix

def get_lp_matrix(n : QKDNetwork):
    if(len(m) == 0):
        for _ in range(100):
            m.append(False)
    nodeNum = len(n.nodeList)
    if(not m[nodeNum]):
        m[nodeNum] = get_lp_solve(n)
    tmp_matrix = deepcopy(m[nodeNum])
    # 随机返回一个业务量矩阵
    for i in range(nodeNum):
        for j in range(nodeNum):
            tmp_matrix[i][j] = random.randint(0, tmp_matrix[i][j])
    return tmp_matrix


def clear_lp_matrix(nodeNum):
    if(len(m) != 0 and m[nodeNum]):
        m[nodeNum] = False
        